import copy
import math

from common_modules import *

def CuptiOverheadFile(cupti_overhead_file):
    t_overhead_cupti = 0.0
    with open(cupti_overhead_file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            t_overhead_cupti += float(row[1][:-3])

    return t_overhead_cupti

def AnalyzerOverheadFile(analyzer_overhead_file):
    t_overhead_analyzer = 0.0
    with open(analyzer_overhead_file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            t_overhead_analyzer += (float(row[1][:-3]) / 1000.0)

    return t_overhead_analyzer

def TimeOverhead(args, devices):
    t_overhead_cupti = {}
    t_overhead_analyzer = {}

    for dev in devices:
        curr_root = args.root[0] + '/' + dev + '/IntraLayer/'
        for net in os.listdir(curr_root):
            if net.lower() == "googlenet":
                continue

            if net not in t_overhead_cupti:
                t_overhead_cupti.setdefault(net, {})
                t_overhead_analyzer.setdefault(net, {})

                if dev not in t_overhead_cupti[net]:
                    t_overhead_cupti[net].setdefault(dev, {})
                    t_overhead_analyzer[net].setdefault(dev, {})

            cupti_overhead_file = curr_root + net + '/CuptiOverhead/CUPTI-OVERHEAD.csv'
            analyzer_overhead_file = curr_root + net + '/AnalyzerOverhead/' + \
                os.listdir(curr_root + net + '/AnalyzerOverhead/')[0]
            # analyzer_overhead_file = curr_root + net + '/AnalyzerOverhead/119-10-23-ANALYZER.csv'

            t_overhead_cupti[net][dev] = CuptiOverheadFile(cupti_overhead_file)
            t_overhead_analyzer[net][dev] = AnalyzerOverheadFile(analyzer_overhead_file)

    t_overhead_folder = args.output[0] + '/TimeOutput'
    if not os.path.exists(args.output[0] + '/TimeOutput'):
        os.makedirs(args.output[0] + '/TimeOutput')

    fields = ['Net Model', 'Device', 'T_p', 'T_a', 'T_total']
    f = open(t_overhead_folder + '/TimeOverhead.csv', 'w')
    writer = csv.DictWriter(f, fieldnames=fields)
    writer.writeheader()

    for net, net_vals in t_overhead_cupti.items():
        for dev, dev_vals in t_overhead_cupti[net].items():
            t_total = t_overhead_cupti[net][dev] + t_overhead_analyzer[net][dev]
            with open(t_overhead_folder + '/TimeOverhead.csv', 'w') as csvfile:
                writer.writerow({fields[0]: net, fields[1]: dev, fields[2]: t_overhead_cupti[net][dev],
                                 fields[3]: t_overhead_analyzer[net][dev], fields[4]: t_total})
